import pyotp
import qrcode
import time
import json
import os
import io
from flask import Flask, render_template, request, jsonify, send_file, flash, redirect, url_for

# 确保 static 目录存在
if not os.path.exists('static'):
    os.makedirs('static')

class TOTPAuthenticator:
    def __init__(self, config_file="accounts.json"):
        self.config_file = config_file
        self.accounts = self._load_accounts()

    def _load_accounts(self):
        if os.path.exists(self.config_file):
            try:
                with open(self.config_file, 'r') as f:
                    return json.load(f)
            except json.JSONDecodeError:
                print(f"警告：{self.config_file} 文件格式错误或为空，将创建新的配置。")
                return {}
        return {}

    def _save_accounts(self):
        with open(self.config_file, 'w') as f:
            json.dump(self.accounts, f, indent=4)

    def add_account(self, account_name):
        """添加新账户并生成密钥"""
        if account_name in self.accounts:
            raise ValueError(f"账户 '{account_name}' 已存在")
        secret = pyotp.random_base32()
        self.accounts[account_name] = secret
        self._save_accounts()
        return secret

    def generate_qrcode_image(self, account_name, issuer="MyAuthApp"):
        """生成二维码图像数据"""
        secret = self.accounts.get(account_name)
        if not secret:
            raise ValueError("Account not found")

        uri = pyotp.totp.TOTP(secret).provisioning_uri(
            name=account_name,
            issuer_name=issuer
        )
        img = qrcode.make(uri)
        
        # 将图片保存到内存中的字节流
        img_byte_arr = io.BytesIO()
        img.save(img_byte_arr, format='PNG')
        img_byte_arr.seek(0) # 重置指针到开头
        return img_byte_arr

    def get_current_code(self, account_name):
        secret = self.accounts.get(account_name)
        if not secret:
            return None, 0

        totp = pyotp.TOTP(secret)
        remaining = 30 - (int(time.time()) % 30)
        return totp.now(), remaining

    def get_all_codes(self):
        """获取所有账户的当前验证码和剩余时间"""
        codes = {}
        for account in self.accounts:
            code, remaining = self.get_current_code(account)
            if code:
                codes[account] = {"code": code, "remaining": remaining}
        return codes

    def verify_code(self, account_name, code):
        """验证用户输入的验证码"""
        secret = self.accounts.get(account_name)
        if not secret:
            return False # 账户不存在
        totp = pyotp.TOTP(secret)
        # 验证当前以及上一个时间窗口的 code，以防时间略有偏差
        return totp.verify(code, valid_window=1)

# --- Flask App --- #

app = Flask(__name__)
app.secret_key = os.urandom(24) # 用于 flash 消息的密钥
auth = TOTPAuthenticator()

@app.route('/')
def index():
    """显示令牌列表页面"""
    return render_template('index.html')

@app.route('/get_tokens')
def get_tokens():
    """提供令牌数据的 API 端点"""
    return jsonify(auth.get_all_codes())

@app.route('/verify', methods=['POST'])
def verify_token():
    """验证提供的令牌是否有效"""
    data = request.get_json()
    if not data or 'account_name' not in data or 'code' not in data:
        return jsonify({'error': 'Missing account_name or code in request'}), 400

    account_name = data['account_name']
    code = data['code']

    is_valid = auth.verify_code(account_name, code)

    if is_valid:
        return jsonify({'status': 'success', 'message': '验证码有效'})
    else:
        # 返回 401 Unauthorized 更符合 RESTful 风格
        return jsonify({'status': 'failure', 'message': '验证码无效或账户不存在'}), 401

@app.route('/add', methods=['GET', 'POST'])
def add_account_route():
    """处理添加账户的页面和逻辑"""
    qr_code_url = None
    if request.method == 'POST':
        account_name = request.form.get('account_name')
        if not account_name:
            flash('账户名称不能为空！', 'error')
        else:
            try:
                auth.add_account(account_name)
                flash(f'账户 "{account_name}" 添加成功！请扫描下方二维码。', 'success')
                # 生成用于模板显示的二维码 URL
                qr_code_url = url_for('serve_qrcode', account_name=account_name, _external=True)
            except ValueError as e:
                flash(str(e), 'error')
            except Exception as e:
                 flash(f'添加账户时发生错误: {e}', 'error')

    # 即使是 POST 请求，如果成功或失败，都重新渲染 add_account.html
    # GET 请求也渲染此页面
    return render_template('add_account.html', qr_code_url=qr_code_url)


@app.route('/qrcode/<account_name>.png')
def serve_qrcode(account_name):
    """提供特定账户的二维码图像"""
    try:
        img_io = auth.generate_qrcode_image(account_name)
        return send_file(img_io, mimetype='image/png')
    except ValueError:
        return "Account not found", 404
    except Exception as e:
        return f"Error generating QR code: {e}", 500

if __name__ == '__main__':
    # 使用 waitress 或 gunicorn 等生产服务器部署
    # 这里为了简单起见，使用 Flask 的开发服务器
    print("Web 服务已启动，请在浏览器中访问 http://127.0.0.1:5000")
    app.run(debug=True) # debug=True 方便开发，生产环境应设为 False
